import pandas as pd
import geopandas as gpd
import matplotlib.pyplot as plt

# 读取CSV数据
data = {
    'NOC': ['USA', 'GBR', 'CHN', 'FRA', 'RUS', 'AUS', 'JPN', 'ITA', 'DEU', 'CAN',
            'NLD', 'KOR', 'BEL', 'BRA', 'NZL', 'HUN', 'ESP', 'ROU', 'UKR', 'NOR',
            'CUB', 'IRN', 'POL', 'KEN', 'UZB', 'DNK', 'SWE', 'TUR', 'CZE', 'AZE',
            ],

    'Predicted Total Medals for 2028': [
        119.7367, 108.58, 88.4767, 63.07, 58.9967, 51.3367, 50.9667, 36.51, 36.16, 32.56,
        28.9367, 28.1367, 26.3733, 19.71, 19.26, 19.14, 18.51, 17.4067, 15.59, 11.9233,
        11.7333, 11.0467, 10.8967, 10.6367, 10.24, 9.9833, 9.98, 9.0033, 8.1467, 8.0233,
    ]
}

df = pd.DataFrame(data)

df['ISO_A3'] = df['NOC']

# 加载世界地图
world = gpd.read_file('path_to_downloaded_file/ne_110m_admin_0_countries.shp')

# 合并数据，确保使用正确的列名
world = world.merge(df, left_on='SOV_A3', right_on='ISO_A3', how='left')

# 剔除南极洲
world = world[world['NAME'] != 'Antarctica']

# 绘制地图
fig, ax = plt.subplots(1, 1, figsize=(15, 10))

# 绘制国家边界
world.boundary.plot(ax=ax, linewidth=1)

# 根据预测的奖牌总数进行填色，使用柔和的颜色填充无数据区域
world.plot(column='Predicted Total Medals for 2028', ax=ax,
           legend_kwds={'label': "2028年预测总奖牌数", 'orientation': "horizontal"},
           cmap='OrRd', linewidth=0.8, edgecolor='0.8', alpha=0.7, legend=True)

# 设置放大显示区域
ax.set_xlim(-180, 180)  # 经度范围
ax.set_ylim(-60, 90)    # 纬度范围

# 设置标题
plt.title('Hot map of countries predicted total medal counts for 2028')

# 设置无数据区域的颜色为柔和的灰色
world['Predicted Total Medals for 2028'].fillna(0, inplace=True)
world.plot(column='Predicted Total Medals for 2028', ax=ax,
           cmap='OrRd', linewidth=0.8, edgecolor='0.8', alpha=0.3)

# 去除坐标轴
ax.set_axis_off()

# 调整颜色条，变细
cbar = plt.colorbar(ax.collections[0], ax=ax, orientation="horizontal", fraction=0.02, pad=0.1)

# 显示地图
plt.show()
